import pandas as pd

def run_backtest(stock_code: str, start_date: str, end_date: str, initial_cash: float):
    """
    Performs a stock backtest based on specific buy/sell signals and pricing rules.

    Args:
        stock_code (str): The stock code to backtest (e.g., "000753").
        start_date (str): The start date for backtesting (YYYY-MM-DD).
        end_date (str): The end date for backtesting (YYYY-MM-DD).
        initial_cash (float): The initial amount of cash for the portfolio.

    Returns:
        dict: A dictionary containing transaction records, portfolio history, and backtest summary.
    """
    # --- Configuration ---
    data_file = 'job/data/trade_indicators.parquet'

    # --- Data Loading and Preprocessing ---
    try:
        df = pd.read_parquet(data_file)
        print(f"Successfully loaded data from '{data_file}'.")
        print("Loaded DataFrame columns:", df.columns)
    except FileNotFoundError:
        print(f"Error: Data file '{data_file}' not found.")
        print("Please ensure 'trade_indicators.parquet' file is present in 'job/data' directory or provide its full path.")
        return {
            "transactions": [],
            "portfolio_history": [],
            "summary": {"error": f"Data file '{data_file}' not found."}
        }
    except Exception as e:
        print(f"Error loading Parquet file: {e}")
        return {
            "transactions": [],
            "portfolio_history": [],
            "summary": {"error": f"Error loading Parquet file: {e}"}
        }

    # Ensure index is not 'date'. If it's already set, reset it first.
    if df.index.name == 'date':
        df = df.reset_index()
        print("DataFrame index reset because its name was 'date'.")

    # Check if 'datekey' column exists, if so, rename to 'date' and convert to datetime objects
    if 'datekey' in df.columns:
        df['date'] = df['datekey'].astype(str).str.replace('d', '') # Remove trailing 'd'
        df['date'] = pd.to_datetime(df['date'], format='%Y%m%d')
        print("Converted 'datekey' column to 'date' column and datetime format.")
    else:
        print("Error: 'datekey' or 'date' column not found in data. Please check the data file.")
        return {
            "transactions": [],
            "portfolio_history": [],
            "summary": {"error": "Missing 'datekey' or 'date' column in data."}
        }

    # Convert 'code' column to string type for consistent filtering
    df['code'] = df['code'].astype(str)

    # Filter data based on specified stock code and date range
    filtered_df = df[(df['code'] == str(stock_code)) &
                     (df['date'] >= pd.to_datetime(start_date)) &
                     (df['date'] <= pd.to_datetime(end_date))].copy()

    # Sort filtered data by date to ensure chronological order for backtesting
    filtered_df = filtered_df.sort_values(by='date').reset_index(drop=True)

    if filtered_df.empty:
        print(f"No data found for stock code '{stock_code}' in date range {start_date} to {end_date}.")
        return {
            "transactions": [],
            "portfolio_history": [],
            "summary": {"message": f"No data found for stock code '{stock_code}' in the specified date range."}
        }

    # --- Prepare signals and pricing data based on strategy rules ---
    if 'action' not in filtered_df.columns:
        print("\nWarning: 'action' column not directly found in data.")
        print("Deriving 'action' column: 'buy' if 'DRAGONBUY' is True, 'sell' if 'POWERDOWNSELL' is True.")
        filtered_df['action'] = None
        filtered_df.loc[filtered_df['DRAGONBUY'] == True, 'action'] = 'buy'
        filtered_df.loc[filtered_df['POWERDOWNSELL'] == True, 'action'] = 'sell'

    filtered_df['next_day_action'] = filtered_df['action'].shift(1)
    filtered_df['ma_5_close'] = filtered_df['close'].rolling(window=5, min_periods=1).mean().shift(1)

    # --- Backtesting Simulation ---
    cash = float(initial_cash)
    shares_held = 0
    transactions = [] # List to store buy/sell transaction records
    portfolio_history = [] # List to store daily portfolio value snapshots

    print(f"\nStarting backtest for stock {stock_code} from {start_date} to {end_date} with initial cash {initial_cash:.2f}.")

    # Iterate through each day in the filtered data to simulate trades
    for i in range(len(filtered_df)):
        current_day_data = filtered_df.iloc[i]
        current_date = current_day_data['date']
        current_close_price = current_day_data['close']
        current_open_price = current_day_data['open']
        current_high_price = current_day_data['high']
        current_low_price = current_day_data['low']
        action_for_today = current_day_data['next_day_action']
        buy_price_ma = current_day_data['ma_5_close']

        # Skip this day if critical price data is missing
        if pd.isna(current_close_price) or pd.isna(current_open_price) or pd.isna(current_high_price) or pd.isna(current_low_price):
            print(f"Warning: Skipping {current_date.strftime('%Y-%m-%d')}: Missing 'close', 'open', 'high', or 'low' price data.")
            continue

        # Record daily portfolio value BEFORE any potential transactions for the day.
        current_portfolio_value_before_trade = cash + (shares_held * current_close_price if shares_held > 0 else 0)
        portfolio_history.append({
            'date': current_date.strftime('%Y-%m-%d'),
            'portfolio_value': round(current_portfolio_value_before_trade, 2),
            'cash': round(cash, 2),
            'shares_held': shares_held,
            'stock_price': round(current_close_price, 2)
        })

        # --- Sell Logic ---
        if action_for_today == 'sell' and shares_held > 0:
            sell_price = current_open_price
            sell_amount = shares_held * sell_price
            cash += sell_amount
            print(f"{current_date.strftime('%Y-%m-%d')}: SOLD {shares_held} shares at {sell_price:.2f}. "
                  f"Proceeds: {sell_amount:.2f}. Current Cash: {cash:.2f}")
            transactions.append({
                'date': current_date.strftime('%Y-%m-%d'),
                'type': 'SELL',
                'stock_code': stock_code,
                'price': round(sell_price, 2),
                'shares': shares_held,
                'amount': round(sell_amount, 2),
                'cash_after_transaction': round(cash, 2),
                'shares_after_transaction': 0
            })
            shares_held = 0

        # --- Buy Logic ---
        elif action_for_today == 'buy' and not pd.isna(buy_price_ma) and buy_price_ma > 0:
            # If 5-day MA is within current day's high and low, use MA for buy price.
            # Otherwise, use current day's open price.
            if current_low_price <= buy_price_ma <= current_high_price:
                buy_price = buy_price_ma
            else:
                buy_price = current_open_price

            if cash >= buy_price:
                shares_to_buy = int(cash / buy_price)
                if shares_to_buy > 0:
                    buy_amount = shares_to_buy * buy_price
                    cash -= buy_amount
                    print(f"{current_date.strftime('%Y-%m-%d')}: BOUGHT {shares_to_buy} shares at {buy_price:.2f}. "
                          f"Cost: {buy_amount:.2f}. Current Cash: {cash:.2f}")
                    transactions.append({
                        'date': current_date.strftime('%Y-%m-%d'),
                        'type': 'BUY',
                        'stock_code': stock_code,
                        'price': round(buy_price, 2),
                        'shares': shares_to_buy,
                        'amount': round(buy_amount, 2),
                        'cash_after_transaction': round(cash, 2),
                        'shares_after_transaction': shares_held + shares_to_buy
                    })
                    shares_held += shares_to_buy

    # --- Final Summary ---
    final_portfolio_value = cash
    if shares_held > 0 and not filtered_df.empty:
        last_close_price = filtered_df.iloc[-1]['close']
        if not pd.isna(last_close_price):
            final_portfolio_value += (shares_held * last_close_price)
            print(f"Last available close price: {last_close_price:.2f}")
        else:
            print("Warning: Could not get latest close price for full portfolio value calculation.")
    
    print(f"\n--- Backtest Summary for Stock {stock_code} ---")
    print(f"Initial Cash: {initial_cash:.2f}")
    print(f"Final Cash: {cash:.2f}")
    print(f"Final Shares Held: {shares_held}")
    print(f"Final Portfolio Value: {final_portfolio_value:.2f}")

    # Convert transactions and portfolio_history to DataFrames
    transactions_df = pd.DataFrame(transactions) if transactions else pd.DataFrame()
    portfolio_history_df = pd.DataFrame(portfolio_history) if portfolio_history else pd.DataFrame()

    # Return the results with DataFrames
    return {
        "transactions": transactions_df.to_dict(orient='records'),
        "portfolio_history": portfolio_history_df.to_dict(orient='records'),
        "summary": {
            "stock_code": stock_code,
            "start_date": start_date,
            "end_date": end_date,
            "initial_cash": round(initial_cash, 2),
            "final_cash": round(cash, 2),
            "final_shares_held": shares_held,
            "final_portfolio_value": round(final_portfolio_value, 2),
        }
    }

# This block allows the script to be run directly for testing.
if __name__ == '__main__':
    stock_code = "603178"       
    start_date = "2023-01-01"       
    end_date = "2025-06-01"       
    
    sample_request_data = {
        "stock_code": stock_code,
        "start_date": start_date,
        "end_date": end_date,
        "initial_cash": 100000
    }

    results = run_backtest(
        stock_code=sample_request_data['stock_code'],
        start_date=sample_request_data['start_date'],
        end_date=sample_request_data['end_date'],
        initial_cash=sample_request_data['initial_cash']
    )
    print("\n--- Backtest Results (Returned) ---")
    print(results)